import mat73
import glob
import numpy as np
from scipy.signal import butter, sosfiltfilt
import tqdm
import os

"""
    mat中的数据
        RSVP
        rest
        senloc
        test
        train

    train中的数据
        chanlocs
        data    
        markers_seq
        markers_target
        nbTrials
        nbTrials_nontarget
        nbTrials_target
        online_acc
        srate
        text_result
        text_to_spell
"""


class Load2022Dataset:
    baseline = [-200, 0]  # in ms
    frame = [0, 600]  # in ms

    def __init__(self, matPath: str|list[str], loadDataImmediately: bool = True):
        """
            Description:
                该类用于读取论文中的mat文件

            Parameters:
                matPath:str
                    mat文件所在的路径
                loadDataImmediately:bool
                    是否立即加载数据
                isFile:bool
                    是否是文件
        """
        if isinstance(matPath, str):
            if os.path.isdir(matPath):
                self.fileArray = glob.glob(os.path.join(matPath, "*.mat"))# [0:20]
                if (len(self.fileArray) < 2):
                    raise Exception("文件数量小于2！")
                # self.fileArray = self.fileArray[0:int(len(self.fileArray) * fileRate)]
            else:
                self.fileArray = [matPath]
        elif isinstance(matPath, list):
            self.fileArray = matPath
        else:
            raise Exception("matPath必须是文件或目录")

        if loadDataImmediately:
            self.loadAllData()

    def loadAllData(self):
        self.targetEEGArray = []
        self.nontargetEEGArray = []
        targetEEG = None
        nontargetEEG = None
        self.dataTarSize = 0
        self.dataNonTargetSize = 0
        for file in tqdm.tqdm(self.fileArray, desc="Load Data"):
            f = mat73.loadmat(file)
            for n_calib in range(len(f['train'])):
                cur_eeg = f['train'][n_calib]
                data = np.asarray(cur_eeg['data'])
                srate = cur_eeg['srate']  # 采样率
                data = self.butter_bandpass_filter(data, 0.5, 10, srate, 4)
                markers = cur_eeg['markers_target']

                targetID = np.where(markers == 1)[0]
                nontargetID = np.where(markers == 2)[0]

                tmp_targetEEG = self.extractEpoch3D(data, targetID, srate, self.baseline, self.frame, False)
                tmp_nontargetEEG = self.extractEpoch3D(data, nontargetID, srate, self.baseline, self.frame, False)
                if n_calib == 0:
                    targetEEG = tmp_targetEEG
                    nontargetEEG = tmp_nontargetEEG
                else:
                    targetEEG = np.dstack((targetEEG, tmp_targetEEG))
                    nontargetEEG = np.dstack((nontargetEEG, tmp_nontargetEEG))
            # 一个数据的大小(32, 614, 840)
            self.targetEEGArray.append(targetEEG)
            self.dataTarSize += np.size(targetEEG, 2)
            self.nontargetEEGArray.append(nontargetEEG)
            self.dataNonTargetSize += np.size(nontargetEEG, 2)

    def extractEpoch3D(self, data, event, srate, baseline, frame, opt_keep_baseline):
        # extract epoch from 2D data into 3D [ch x time x trial]
        # time才是数据
        # input: event, baseline, frame
        # extract epoch = baseline[0] to frame[2]

        # for memory pre-allocation
        if opt_keep_baseline == True:
            begin_tmp = int(np.floor(baseline[0] / 1000 * srate))
            end_tmp = int(begin_tmp + np.floor(frame[1] - baseline[0]) / 1000 * srate)
        else:
            begin_tmp = int(np.floor(frame[0] / 1000 * srate))
            end_tmp = int(begin_tmp + np.floor(frame[1] - frame[0]) / 1000 * srate)
            # end_tmp：307，begin_tmp：0
            # scrate：512
            # print(begin_tmp, end_tmp)
        # shape为信号、信号、第几次
        epoch3D = np.zeros((data.shape[0], end_tmp - begin_tmp, len(event)))
        nth_event = 0

        for i in event:
            if opt_keep_baseline == True:
                begin_id = int(i + np.floor(baseline[0] / 1000 * srate))
                end_id = int(begin_id + np.floor((frame[1] - baseline[0]) / 1000 * srate))
            else:
                # 选择frame[0]和frame[1]的时间段
                begin_id = int(i + np.floor(frame[0] / 1000 * srate))
                end_id = int(begin_id + np.floor((frame[1] - frame[0]) / 1000 * srate))

            tmp_data = data[:, begin_id:end_id]

            # 选取baseline的时间段
            begin_base = int(np.floor(baseline[0] / 1000 * srate))
            end_base = int(begin_base + np.floor(np.diff(baseline) / 1000 * srate) - 1)
            base = np.mean(tmp_data[:, begin_base:end_base], axis=1)

            rmbase_data = tmp_data - base[:, np.newaxis]
            epoch3D[:, :, nth_event] = rmbase_data
            nth_event = nth_event + 1

        return epoch3D

    def butter_bandpass_filter(self, data, lowcut, highcut, fs, order):
        nyq = fs / 2
        low = lowcut / nyq
        high = highcut / nyq
        sos = butter(order, [low, high], btype='band', output='sos')
        # demean before filtering
        meandat = np.mean(data, axis=1)
        data = data - meandat[:, np.newaxis]
        y = sosfiltfilt(sos, data)  # zero-phase filter # data: [ch x time]
        # specify pandlen to make the result the same as Matlab filtfilt()
        return y

    def getData(self):
        return self.targetEEGArray, self.nontargetEEGArray

    def getDataNdarray(self):
        return np.array(self.targetEEGArray), np.array(self.nontargetEEGArray)

    def getDataSize(self):
        return self.dataTarSize, self.dataNonTargetSize

    def getDataArraySize(self):
        return len(self.targetEEGArray), len(self.nontargetEEGArray)


if __name__ == "__main__":
    a = Load2022Dataset(r"F:\DataSet\P300_2022Dataset", True)
    target, _ = a.getData()
    print(np.array(target).shape)   #   (2, 32, 307, 300)
    print(np.array(target[0]).shape)#   (32, 307, 300)
    # file = mat73.loadmat(r"E:\文件\Documents\Python\大创\log\s01.mat")
    #
    # a = file["train"]  # len=2
    # for key in a[1].keys():
    #     print(key)
    # a1 = np.where(a[0]["markers_target"] == 1)[0]
    # a2 = np.nonzero(a[1]["markers_target"] == 1)[0]
    # print(len(a1))
    # print(len(a2))
    # print(np.where(a[0]["markers_target"]==1)[0])
    # print(a[0]["nbTrials"])#900.0
    # print(a[0]["nbTrials_nontarget"])  # 750.0
    # print(a[1]["text_to_spell"])
